#! /usr/bin/python
import sys,getopt
import glob
import os.path
import cPickle as pickle

def get_abs_file_path(path):
    return os.path.realpath(os.path.expanduser(path))

class Tagger:
    def __init__(self):
        self.dict_tag_file = {}
        self.dict_file_tag = {}

    def add_tag_to_file(self, tag, file):
        if self.dict_tag_file.has_key(tag):
            self.dict_tag_file[tag].add(file)
        else:
            self.dict_tag_file[tag] = set([file])


        if self.dict_file_tag.has_key(file):
            self.dict_file_tag[file].add(tag)
        else:
            self.dict_file_tag[file] = set([tag])

    def add_tags_to_files(self, tags, files):
        for tag in tags:
            for f in files:
                self.add_tag_to_file(tag,f)

    def remove_tag_from_file(self, tag, file):
        if self.dict_tag_file.has_key(tag):
            self.dict_tag_file[tag].discard(file)
        if self.dict_file_tag.has_key(file):
            self.dict_file_tag[file].discard(tag)
            if len(self.dict_file_tag[file]) == 0:
                del self.dict_file_tag[file]

    def remove_tags_from_files(self,tags,files):
        for tag in tags:
            for f in files:
                self.remove_tag_from_file(tag,f)
            
    def list_files_with_tag(self, tag):
        if self.dict_tag_file.has_key(tag):
            return self.dict_tag_file[tag]
        else:
            return None

    def list_files_with_all_tags(self, list_tags):
        list_tags = [tag for tag in list_tags if self.dict_tag_file.has_key(tag)]
        if len(list_tags) == 0:
            return set([])
        ret_set = self.dict_tag_file[list_tags[0]]
        for tag in list_tags:
            ret_set &= self.dict_tag_file[tag]
        return ret_set

    def list_all_tags(self):
        return dict_tag_file.keys()

    def reset(self):
        self.dict_tag_file = {}
        self.dict_file_tag = {}

    def dump_state(self):
        print "dict_tag_file", self.dict_tag_file
        print "dict_file_tag", self.dict_file_tag

class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg

class CommandEnum:
    HELP = 1
    ASSIGN = 2
    LIST_TAGS = 3
    LIST_FILES = 4
    UNASSIGN = 5
    DUMP_STATE = 6
    RESET = 7
        
class Arguments:
    def __init__(self):
        self.command = None
        self.tags = None
        self.files = None
        self.database = "~/.tagit.pkl"

    def parse_from_arguments(self, list_args):
        # Parse all the commandline options
        try:
            opts, args = getopt.getopt(list_args, ":ht:d:", ["help"])
        except getopt.error, msg:
            raise Usage(msg)
        for o, v in opts:
            if o == "-t":
                self.tags = [f for f in v.split(',') if len(f)>0]
                if len(self.tags) == 0:
                    raise Usage("No tags specified")
            elif o in ("-h", "--help"):
                self.command = CommandEnum.HELP
            elif o == "-d":
                if os.path.isdir(get_abs_file_path(v)): raise Usage("Given path is a directory : " + v)
                self.database = v
            else:
                raise Usage("Unrecognized option "+o)

        if self.command == CommandEnum.HELP:
            return
        if len(args) == 0:
            raise Usage("Command not specified")
        else:
            if args[0] == "a": self.command = CommandEnum.ASSIGN
            elif args[0] == "lt": self.command = CommandEnum.LIST_TAGS
            elif args[0] == "lf": self.command = CommandEnum.LIST_FILES
            elif args[0] == "r": self.command = CommandEnum.UNASSIGN
            elif args[0] == "dump": self.command = CommandEnum.DUMP_STATE
            elif args[0] == "reset": self.command = CommandEnum.RESET
            else: raise Usage("Command not recognized : " + args[0])

            if len(args) > 1:
                set_files = set([])
                for arg in args[1:]:
                    set_files.update(glob.glob(arg))
                self.files = set([get_abs_file_path(f) for f in set_files])
        self.database = get_abs_file_path(self.database)
        self.validate_args()

    def validate_args(self):
        if self.command == CommandEnum.ASSIGN and (self.files is None or self.tags is None):
            raise Usage("Both files and tags must be specified for assignment")
        elif self.command == CommandEnum.UNASSIGN and (self.files is None or self.tags is None):
            raise Usage("Both files and tags must be specified for un-assignment")

    def dump_str(self):
        print "command : ", self.command
        print "database : ", self.database
        print "tags : ", self.tags
        print "files : ", self.files

        

def print_usage():
    print """USAGE::
    tagit [-h] [--help] [-t TAGS] [-d DATA_FILE] [COMMAND] [FILES]
    OPTIONS are:
    -h, --help: print this usage information
    -t        : comma separated list of tags
    -d        : database file (defaults to ~/.tagit.pkl)

    COMMAND: Available commands are
    a         : Assign specified tags to specified files
    r         : Remove specified tags from specified files
    lt        : list tags
    lf        : list files corresponding to the given tags

    FILES     : can be any list of files separated by spaces. shell glob patterns are also accepted. This is a required argument unless the -l or -h flags are specified"""

def print_set(set_files):
    for f in set_files: print f
        
        
def main(argv=None):
    if argv is None:
        argv = sys.argv

    try:
        # Parse all the commandline options
            
        args = Arguments()
        args.parse_from_arguments(argv[1:])
        
        # Parsing complete
        # Now we unpickle

        if not (os.path.exists(args.database) and os.path.isfile(args.database)):
            tgr = Tagger()
        else:
            pf = open(args.database, "r")
            tgr = pickle.load(pf)
            pf.close()

        # Here the actual work is done.
        is_modified = False
            
        if args.command == CommandEnum.LIST_FILES:
            print_set(tgr.list_files_with_all_tags(args.tags))
        elif args.command == CommandEnum.UNASSIGN:
            tgr.remove_tags_from_files(args.tags, args.files)
            is_modified = True
        elif args.command == CommandEnum.ASSIGN:
            tgr.add_tags_to_files(args.tags, args.files)
            is_modified = True
        elif args.command == CommandEnum.LIST_TAGS:
            pass
        elif args.command == CommandEnum.HELP:
            print_usage()
        elif args.command == CommandEnum.DUMP_STATE:
            tgr.dump_state()
        elif args.command == CommandEnum.RESET:
            tgr.reset()
            is_modified = True
        else: raise Usage("Unknown command : " + args.command)
            
        # Finally pickle the tagger.
        if is_modified:
            wpf = open(args.database, "w")
            pickle.dump(tgr, wpf, 1)
            wpf.close()
         
    except Usage, err:
        print >>sys.stderr, err.msg
        print_usage()
        return 2

if __name__ == "__main__":
    sys.exit(main())
